package main

import (
	"encoding/csv"
	"lenet/network"
	"lenet/network/activations"
	"lenet/network/layers"
	"log"
	"os"
	"strconv"

	"gonum.org/v1/gonum/floats"
)

func main() {
	lRate := 0.00000000001
	leNet5 := network.Network{}

	leNet5.AddLayer(layers.NewPaddingLayer(2))
	leNet5.AddLayer(layers.NewConvLayer(1, 32, 5, 6, 1, 1, lRate, activations.Relu{}))

	leNet5.AddLayer(layers.NewMaxPull2d(6, 28, 2, 2))
	leNet5.AddLayer(layers.NewConvLayer(6, 14, 5, 16, 1, 1, lRate, activations.Relu{}))
	leNet5.AddLayer(layers.NewMaxPull2d(16, 10, 2, 2))

	leNet5.AddLayer(layers.NewFullLayer(5*5*16, 120, lRate, activations.Relu{}))
	leNet5.AddLayer(layers.NewFullLayer(120, 84, lRate, activations.Relu{}))
	leNet5.AddLayer(layers.NewFullLayer(84, 10, lRate, activations.Linear{}))

	inputs, results, err := readFile("../mnist_train.csv")
	if err != nil {
		log.Fatal(err)
	}

	lenT := 0
	//обучение
	for i := 0; i < len(inputs); i++ {
		if i%1000 == 0 {
			log.Print("lening ", i, " ")
		}
		res := leNet5.Forward(inputs[i])

		tmp := make([]float64, 10)
		tmp[int(results[i])] = 10
		//log.Println(floats.MaxIdx(res), int(results[i]), res, floats.MaxIdx(res) == int(results[i]))
		if floats.MaxIdx(res) == int(results[i]) {
			lenT++
		}

		if i%1000 == 0 {
			log.Print("Acc ", i, ": ", float64(lenT)/1000.0, lenT)
			log.Print(res)
			lenT = 0.
		}

		floats.Sub(tmp, res)

		leNet5.BackProp(tmp)
		//var a int
		//fmt.Scan(&a)
	}

	inputs, results, err = readFile("../mnist_test.csv")
	if err != nil {
		log.Fatal(err)
	}

	//тест
	trueRes := 0
	falseRes := 0
	for i := 0; i < len(inputs); i++ {
		log.Println("test", i)
		res := leNet5.Forward(inputs[i])

		resNum := floats.MaxIdx(res)

		if resNum == int(results[i]) {
			trueRes++
		} else {
			falseRes++
		}
	}

	log.Println("True:", trueRes)
	log.Println("False:", falseRes)
	log.Println("Accuracy:", float64(trueRes)/float64((trueRes+falseRes)))

}

func readFile(path string) ([][]float64, []float64, error) {
	file, err := os.Open(path)
	if err != nil {
		log.Println(err)
		return nil, nil, err
	}
	defer file.Close()

	reader := csv.NewReader(file)
	records, err := reader.ReadAll()
	if err != nil {
		log.Println(err)
		return nil, nil, err
	}

	//del names
	records = records[1:]

	floatRecords := make([][]float64, 0)
	for i, record := range records {
		floatRecord := make([]float64, 28*28+1)
		floatRecords = append(floatRecords, floatRecord)
		for j := 0; j < len(record); j++ {
			f, err := strconv.Atoi(record[j])
			if err != nil {
				log.Println(err)
				return nil, nil, err
			}
			floatRecords[i][j] = float64(f)
		}
	}

	var results []float64
	var inputs [][]float64
	for _, record := range floatRecords {
		results = append(results, record[0])
		inputs = append(inputs, record[1:])
	}

	return inputs, results, nil
}
